#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
从conda环境导出适合pip安装的requirements.txt文件
"""

import os
import sys
import subprocess
import argparse
import re
from pathlib import Path

# conda专用包列表，这些包不会被包含在pip requirements中
CONDA_SPECIFIC_PACKAGES = [
    'bzip2', 'ca-certificates', 'libarchive', 'libcurl', 'libffi', 'libiconv',
    'libmamba', 'libsolv', 'libssh2', 'libxml2', 'lz4-c', 'openssl', 'pcre2',
    'tk', 'xz', 'yaml-cpp', 'zlib', 'zstandard', 'zstd', 'vc', 'vs2015_runtime',
    'conda', 'conda-package-handling', 'conda-package-streaming', 'menuinst',
    'pycosat', 'ruamel.yaml', 'win-inet-pton', 'anaconda-anon-usage', 'anaconda_powershell_prompt',
    'anaconda_prompt', 'libmambapy', 'cpp-expected', 'fmt', 'nlohmann_json', 'pybind11-abi',
    'reproc', 'reproc-cpp', 'simdjson', 'spdlog', 'yaml-cpp', 'brotli-python', 'python', 'sqlite',
    'ucrt', 'vc14_runtime', 'tzdata'
]

# 包名称映射，将conda包名映射到pip包名
PACKAGE_NAME_MAPPING = {
    'brotli-python': 'Brotli',
    'typing_extensions': 'typing-extensions',
}

# 特殊版本处理，将特殊版本替换为兼容版本
VERSION_REPLACEMENTS = {
    r'^\d{4}[a-z]$': lambda v: v[:-1],  # 例如：2025a -> 2025
    r'\+cu\d+$': lambda v: '',          # 例如：2.5.1+cu121 -> 2.5.1
    r'\+cpu$': lambda v: '',            # 例如：2.5.1+cpu -> 2.5.1
}

def list_conda_environments():
    """列出所有conda环境"""
    result = subprocess.run(["conda", "env", "list"], capture_output=True, text=True)
    
    if result.returncode != 0:
        print(f"获取conda环境列表失败: {result.stderr}")
        return []
    
    envs = []
    for line in result.stdout.split('\n'):
        if line.startswith('#') or not line.strip():
            continue
        
        parts = line.split()
        if len(parts) >= 1:
            envs.append(parts[0])
    
    return envs

def normalize_version(version):
    """标准化版本号，处理特殊版本"""
    for pattern, replacement_func in VERSION_REPLACEMENTS.items():
        if re.search(pattern, version):
            return re.sub(pattern, replacement_func(version), version)
    return version

def export_pip_requirements(env_name, output_path=None):
    """导出适合pip安装的requirements.txt文件"""
    if not output_path:
        output_path = f"{env_name}_pip_requirements.txt"
    
    # 检查环境是否存在
    envs = list_conda_environments()
    if env_name not in envs:
        print(f"错误: 环境 {env_name} 不存在")
        return False
    
    print(f"正在从环境 {env_name} 导出pip包列表...")
    
    # 获取环境中的所有包
    result = subprocess.run(["conda", "list", "--name", env_name], capture_output=True, text=True)
    
    if result.returncode != 0:
        print(f"获取环境包列表失败: {result.stderr}")
        return False
    
    # 解析包列表，排除conda专用包
    pip_packages = {}  # 使用字典避免重复
    for line in result.stdout.split('\n'):
        if line.startswith('#') or not line.strip():
            continue
        
        parts = line.split()
        if len(parts) >= 2:
            package_name = parts[0].lower()
            version = parts[1]
            
            # 排除conda专用包
            if package_name not in CONDA_SPECIFIC_PACKAGES:
                # 处理包名映射
                if package_name in PACKAGE_NAME_MAPPING:
                    package_name = PACKAGE_NAME_MAPPING[package_name]
                
                # 标准化版本号
                version = normalize_version(version)
                
                # 使用字典存储，避免重复
                pip_packages[package_name] = version
    
    # 写入requirements.txt文件
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write("# 从conda环境 {} 导出的pip requirements\n".format(env_name))
        f.write("# 此文件适用于 pip install -r 命令\n\n")
        
        for package_name, version in sorted(pip_packages.items()):
            f.write(f"{package_name}=={version}\n")
    
    print(f"已成功导出pip requirements到文件: {output_path}")
    print(f"共导出 {len(pip_packages)} 个包")
    return True

def main():
    parser = argparse.ArgumentParser(description="从conda环境导出适合pip安装的requirements.txt文件")
    parser.add_argument("--env-name", type=str, help="要导出的conda环境名称")
    parser.add_argument("--output", type=str, help="输出文件路径")
    parser.add_argument("--list", action="store_true", help="列出所有conda环境")
    
    args = parser.parse_args()
    
    if args.list:
        print("可用的conda环境:")
        for env in list_conda_environments():
            print(f"- {env}")
        return
    
    if not args.env_name:
        # 如果没有指定环境名称，列出可用选项并让用户选择
        available_envs = list_conda_environments()
        if not available_envs:
            print("没有找到任何conda环境")
            return
        
        print("请选择要导出的conda环境:")
        for i, env in enumerate(available_envs):
            print(f"{i+1}. {env}")
        
        try:
            choice = int(input("请输入选项编号: "))
            if 1 <= choice <= len(available_envs):
                env_name = available_envs[choice-1]
            else:
                print("无效的选择")
                return
        except ValueError:
            print("请输入有效的数字")
            return
    else:
        env_name = args.env_name
    
    export_pip_requirements(env_name, args.output)

if __name__ == "__main__":
    main() 